package rsax

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"hash"

	"gitcode.com/yeyixianyang/gokits/encoder"
)

// Encrypt 使用公钥加密
func Encrypt(key *rsa.PublicKey, plainText []byte) ([]byte, error) {
	partLen := key.Size() - 11
	chunks := split(plainText, partLen)

	buffer := bytes.NewBufferString("")

	for _, chunk := range chunks {
		bs, err := rsa.EncryptPKCS1v15(rand.Reader, key, chunk)
		if err != nil {
			return nil, err
		}

		buffer.Write(bs)
	}

	return buffer.Bytes(), nil
}

// EncryptWithEncoder 使用公钥加密，返回加密后的、经过编码的字符串，支持base64、十六进制及二进制编码
func EncryptWithEncoder(key *rsa.PublicKey, plainText []byte, enc encoder.Encoder) (string, error) {
	cipherText, err := Encrypt(key, plainText)
	if err != nil {
		return "", err
	}

	return enc.EncodeToString(cipherText), nil
}

// Decrypt 使用私钥解密.
func Decrypt(key *rsa.PrivateKey, cipherText []byte) (plainText []byte, err error) {
	chunks := split(cipherText, key.PublicKey.Size())

	buffer := bytes.NewBufferString("")

	for _, chunk := range chunks {
		decrypted, err := rsa.DecryptPKCS1v15(rand.Reader, key, chunk)
		if err != nil {
			return nil, err
		}

		buffer.Write(decrypted)
	}

	return buffer.Bytes(), err
}

// DecryptWithEncoder 使用私钥解密，返回解密后的字符串，支持经过base64、十六进制及二进制编码的字符串解密.
func DecryptWithEncoder(key *rsa.PrivateKey, cipherText string, enc encoder.Encoder) ([]byte, error) {
	cipher, err := enc.DecodeString(cipherText)
	if err != nil {
		return nil, err
	}

	plainText, err := Decrypt(key, cipher)
	if err != nil {
		return nil, err
	}

	return plainText, nil
}

// EncryptOAEPWithSha256 使用SHA256的OAEP公钥加密.
func EncryptOAEPWithSha256(key *rsa.PublicKey, plainText []byte) ([]byte, error) {
	return encryptOAEP(key, plainText, sha256.New())
}

// EncryptOAEPWithSha256AndEncoder 使用SHA256的OAEP公钥加密，返回加密后的、经过编码的字符串，支持base64、十六进制及二进制编码.
func EncryptOAEPWithSha256AndEncoder(key *rsa.PublicKey, plainText []byte, enc encoder.Encoder) (string, error) {
	cipher, err := EncryptOAEPWithSha256(key, plainText)
	if err != nil {
		return "", err
	}

	return enc.EncodeToString(cipher), nil
}

// DecryptOAEPWithSha256 使用SHA256的OAEP私钥解密.
func DecryptOAEPWithSha256(key *rsa.PrivateKey, cipherText []byte) ([]byte, error) {
	return decryptOAEP(key, cipherText, sha256.New())
}

// DecryptOAEPWithSha256AndEncoder 使用SHA256的OAEP私钥解密，返回解密后的字符串，支持经过base64、十六进制及二进制编码的字符串解密.
func DecryptOAEPWithSha256AndEncoder(key *rsa.PrivateKey, cipherText string, enc encoder.Encoder) ([]byte, error) {
	cipher, err := enc.DecodeString(cipherText)
	if err != nil {
		return nil, err
	}

	return DecryptOAEPWithSha256(key, cipher)
}

// EncryptOAEPWithSha1 使用SHA1的OAEP公钥加密.
func EncryptOAEPWithSha1(key *rsa.PublicKey, plainText []byte) ([]byte, error) {
	return encryptOAEP(key, plainText, sha1.New())
}

// EncryptOAEPWithSha1AndEncoder 使用SHA1的OAEP公钥加密，返回加密后的、经过编码的字符串，支持base64、十六进制及二进制编码
func EncryptOAEPWithSha1AndEncoder(key *rsa.PublicKey, plainText []byte, enc encoder.Encoder) (string, error) {
	cipherText, err := EncryptOAEPWithSha1(key, plainText)
	if err != nil {
		return "", err
	}

	return enc.EncodeToString(cipherText), nil
}

// DecryptOAEPWithSha1 使用SHA1的OAEP私钥解密
func DecryptOAEPWithSha1(key *rsa.PrivateKey, cipherText []byte) ([]byte, error) {
	return decryptOAEP(key, cipherText, sha1.New())
}

// DecryptOAEPWithSha1AndEncoder 使用SHA1的OAEP私钥解密，返回解密后的字符串，支持经过base64、十六进制及二进制编码的字符串解密
func DecryptOAEPWithSha1AndEncoder(key *rsa.PrivateKey, cipherText string, enc encoder.Encoder) ([]byte, error) {
	cipher, err := enc.DecodeString(cipherText)
	if err != nil {
		return nil, err
	}

	return DecryptOAEPWithSha1(key, cipher)
}

func encryptOAEP(key *rsa.PublicKey, plainText []byte, hash hash.Hash) ([]byte, error) {
	partLen := key.Size() - 2*hash.Size() - 2
	chunks := split(plainText, partLen)

	buffer := bytes.NewBufferString("")

	for _, chunk := range chunks {
		bs, err := rsa.EncryptOAEP(hash, rand.Reader, key, chunk, nil)
		if err != nil {
			return nil, err
		}

		buffer.Write(bs)
	}

	return buffer.Bytes(), nil
}

func decryptOAEP(key *rsa.PrivateKey, cipherText []byte, hash hash.Hash) ([]byte, error) {
	chunks := split(cipherText, key.PublicKey.Size())

	buffer := bytes.NewBufferString("")

	for _, chunk := range chunks {
		decrypted, err := rsa.DecryptOAEP(hash, rand.Reader, key, chunk, nil)
		if err != nil {
			return nil, err
		}

		buffer.Write(decrypted)
	}

	return buffer.Bytes(), nil
}

func split(data []byte, blockSize int) [][]byte {
	var chunk []byte

	chunks := make([][]byte, 0, len(data)/blockSize+1)

	for len(data) >= blockSize {
		chunk, data = data[:blockSize], data[blockSize:]
		chunks = append(chunks, chunk)
	}

	if len(data) > 0 {
		chunks = append(chunks, data)
	}

	return chunks
}
